#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Date : 2024/7/16
# @Author : shiyou pan

from dataclasses import dataclass
from typing import Literal

import pandas as pd

from .date_calc import intnx

# 日期格式
dict_fmt: dict = {'y': '%Y', 'q': '%Y-%m', 'm': '%Y-%m', 'w': '%Y-w%U', 'd': '%Y-%m-%d'}


# 日期序列
def _date_range(min_date: str,
                max_date: str,
                seq_type: Literal['y', 'q', 'm', 'w', 'd']) -> pd.date_range:
    dict_frq: dict = {'y': 'YE', 'q': 'QE-MAR', 'm': 'ME', 'w': 'W', 'd': 'D'}
    date_index = pd.date_range(start=intnx(min_date, interval=seq_type, n=0, sign='b'),
                               end=intnx(max_date, interval=seq_type, n=0, sign='e'),
                               freq=dict_frq.get(seq_type)).strftime(dict_fmt.get(seq_type))
    return date_index


def fill_msdate_wo_key(indata: pd.DataFrame,
                       seq_var: str | None = None,
                       seq_type: Literal['y', 'q', 'm', 'w', 'd'] = 'm') -> pd.DataFrame:
    # 简单设置数据的索引，并进行简单的数据清洗
    data_copy = indata.copy()
    if seq_var:
        data_copy['new_seq'] = data_copy[seq_var].astype(str).apply(
            lambda x: intnx(x, interval=seq_type, n=0, sign='e'))
    else:
        data_copy['new_seq'] = data_copy.index.astype(str).map(lambda x: intnx(x, interval=seq_type, n=0, sign='e'))
    data_copy['new_seq'] = pd.to_datetime(data_copy['new_seq']).dt.strftime(dict_fmt.get(seq_type))
    data_copy.sort_values(by=[seq_var], inplace=True)
    data_copy.drop_duplicates(subset=['new_seq'], keep='last', inplace=True)

    # 填补缺失日期
    min_date = data_copy[seq_var].astype('str').min()
    max_date = data_copy[seq_var].astype('str').max()
    date_index = _date_range(min_date, max_date, seq_type)

    # 插入缺失的日期
    data_copy.set_index(['new_seq'], inplace=True)
    data_copy = data_copy.reindex(date_index)

    return data_copy


def fill_msdate_w_key(indata: pd.DataFrame,
                      key_var: str,
                      seq_var: str,
                      seq_type: Literal['y', 'q', 'm', 'w', 'd'] = 'm') -> pd.DataFrame:
    # 简单设置数据的索引，并进行简单的数据清洗
    data_copy = indata.copy()
    data_copy['new_seq'] = data_copy[seq_var].astype(str).apply(lambda x: intnx(x, interval=seq_type, n=0, sign='e'))
    data_copy['new_seq'] = pd.to_datetime(data_copy['new_seq']).dt.strftime(dict_fmt.get(seq_type))
    data_copy.sort_values(by=[key_var, seq_var], inplace=True)
    data_copy.drop_duplicates(subset=[key_var, 'new_seq'], keep='last', inplace=True)

    # 填补缺失日期
    min_date = data_copy[seq_var].astype('str').min()
    max_date = data_copy[seq_var].astype('str').max()
    date_index = _date_range(min_date, max_date, seq_type)

    # 分层的日期区间
    dict_min: dict = data_copy.groupby(key_var)['new_seq'].min().to_dict()
    dict_max: dict = data_copy.groupby(key_var)['new_seq'].max().to_dict()

    # 分层插入缺失的日期
    multi_index = pd.MultiIndex.from_product([data_copy[key_var].unique(), date_index], names=[key_var, 'new_seq'])
    data_copy.set_index([key_var, 'new_seq'], inplace=True)
    data_copy = data_copy.reindex(multi_index)

    # 剔除不需要的行
    data_copy = data_copy[(data_copy.index.get_level_values(1) >= data_copy.index.get_level_values(0).map(dict_min)) &
                          (data_copy.index.get_level_values(1) <= data_copy.index.get_level_values(0).map(dict_max))]
    data_copy.reset_index(inplace=True)

    return data_copy


@dataclass
class TsProcessor:
    indata: pd.DataFrame
    date_col: str | int | None = None
    freq: Literal['y', 'q', 'm', 'w', 'd'] | None = None
    keys: list[str] | None = None
    is_fill_msdate: bool = True

    def __post_init__(self):
        if self.date_col is None:
            pass
        elif isinstance(self.date_col, str):
            self.indata.set_index(self.date_col, inplace=True)
        elif isinstance(self.date_col, int):
            self.indata.set_index(self.indata.columns[self.date_col], inplace=True)
        else:
            raise ValueError('索引列类型错误')


if __name__ == '__main__':
    pass
